package ru.tusur.asu.pnn.perceptron;

import static ru.tusur.asu.utils.MathUtils.ceil4;
import static ru.tusur.asu.utils.MathUtils.dot;
import static ru.tusur.asu.utils.MathUtils.clamp;

import java.util.LinkedList;

import org.apache.commons.lang.Validate;

public class MultilayerPerceptron {
	
	public static class Builder {
		private int mInputCount;
		private final LinkedList<Layer> mLayers = new LinkedList<Layer>();
		
		public Builder setInputCount(int count) {
			mInputCount = count;
			return this;
		}
		
		public Builder addLayer(Layer layer) {
			mLayers.add(layer);
			return this;
		}
		
		public MultilayerPerceptron create() {
			Validate.isTrue(mInputCount > 0, "Input count should be positive");
			
			MultilayerPerceptron result = new MultilayerPerceptron();
			
			// calculate input count
			int totalInputs = mInputCount;
			for (Layer l : mLayers) {
				totalInputs = ceil4(totalInputs + 1) + l.size();	
			}
			
			// allocate data
			result.data = new float[totalInputs];
			
			// fill perceptron topology
			int weightIndex = 0;
			
			int inputIndex = mInputCount;
			int layerStart = 0;
			
			final LinkedList<LayerStruct> layers = new LinkedList<LayerStruct>();
			final LinkedList<NeuronStruct> neurons = new LinkedList<NeuronStruct>();
			
			for (int i = 0; i < mLayers.size(); i++) {
				inputIndex = ceil4(inputIndex + 1);
				
				int layerInputCount = (i == 0) ? mInputCount + 1 : mLayers.get(i - 1).size() + 1;
				final int layerLength = mLayers.get(i).size();
				
				// set layer
				result.data[layerStart + layerInputCount - 1] = 1; // set offset input
				layers.add(new LayerStruct() {{
					startNeuronIndex = neurons.size();
					endNeuronIndex = startNeuronIndex + layerLength;
				}});
				
				for (ActivationFunction func : mLayers.get(i)) {
					NeuronStruct neuron = new NeuronStruct();
					neuron.weightStartIndex = weightIndex;
					neuron.inputStartIndex = layerStart;
					neuron.index = neurons.size();
					neuron.weightCount = layerInputCount;
					neuron.outputIndex = inputIndex++;
					neuron.function = func;
					
					neurons.add(neuron);
					
					weightIndex += ceil4(layerInputCount);
				}
				layerStart += ceil4(layerInputCount);
			}
			
			result.inputCount = mInputCount;
			result.layers = layers.toArray(new LayerStruct[layers.size()]);
			result.neurons = neurons.toArray(new NeuronStruct[neurons.size()]);			
			result.weights = new float[weightIndex];
			result.neuronSumms = new float[neurons.size()];			
			
			return result;
			
		}
	}
	
	private static class LayerStruct {
		public int startNeuronIndex;
		public int endNeuronIndex;
		
		public int size() {
			return endNeuronIndex - startNeuronIndex;
		}
	}
	
	private static class NeuronStruct {
		
		public static final float MAX_INPUT = 150;
		
		public int weightStartIndex;
		public int inputStartIndex;
		public int weightCount;
		public int outputIndex;		 
		public int index; // TODO remove?
		public ActivationFunction function;
	}
	
	private int inputCount;
	private float[] weights;
	private NeuronStruct[] neurons;
	private LayerStruct[] layers;
	private float[] data;
	private float[] neuronSumms;	

	public int getOutputCount() {
		return layers[layers.length - 1].size();
	}	
	
	public void copyOutputTo(float[] output) {
		int outputCount = getOutputCount();	
		System.arraycopy(data, data.length - outputCount, output, 0, outputCount);
	}
	
	public int getInputCount() {
		return inputCount;
	}
	
	public int getParameterCount() {
		return weights.length;
	}
	
	public void copyParametersTo(float[] parameters) {
		System.arraycopy(weights, 0, parameters, 0, weights.length);
	}
	
	public void setParameters(float[] parameters) {
		System.arraycopy(parameters, 0, weights, 0, weights.length);
	}
	
	public void model(float[] input) {
		Validate.isTrue(input != null && input.length == inputCount, 
				"input length should be " + inputCount);
		
		System.arraycopy(input, 0, data, 0, inputCount);
		for (LayerStruct layer : layers) {
			// TODO parallel for neurons in layer
			for (int i = layer.startNeuronIndex; i < layer.endNeuronIndex; i++) {
				activateNeuron(neurons[i]);
			}
		}		
	}
		
	private MultilayerPerceptron() {		
	}
	
	private void activateNeuron(NeuronStruct n) {
		float sum = dot(data, n.inputStartIndex, weights, n.weightStartIndex, n.weightCount);
		sum = clamp(sum, -NeuronStruct.MAX_INPUT, NeuronStruct.MAX_INPUT);
		neuronSumms[n.index] = sum;
		data[n.outputIndex] = n.function.value(sum);
	}
}
